import os
import argparse
import pandas as pd

def get_aishell3_audio_name_text_df(data_root_path,audio_file_column_name="audio_file",text_column_name="text"):

    audio_file_list = []
    text_list = []

    train_content = open(os.path.join(
        data_root_path, "train", "content.txt")).readlines()
    test_content = open(os.path.join(
        data_root_path, "test", "content.txt")).readlines()

    for line in train_content + test_content:
        audio_name, content = line.split("\t")

        text = "".join(content.split(" ")[0::2])

        audio_file_list.append(audio_name)
        text_list.append(text)
    
    df = pd.DataFrame({
        audio_file_column_name:audio_file_list,
        text_column_name:text_list
    })

    return df

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate wav_file - text pair for MFA")
    parser.add_argument("data_root_path", type=str, help="aishell3 data root ")
    parser.add_argument("output_path", type=str, help="aishell3 data root ")
    
    args = parser.parse_args()
    df = get_aishell3_audio_name_text_df(args.data_root_path)
    if args.output_path.endswith(".csv"):

        df.to_csv(args.output_path,index=False)

    elif args.output_path.endswith(".xlsx") or args.output_path.endswith(".xls"):

        df.to_excel(args.output_path,index=False)
    
    print("========= output head ===============")
    print(df.head())
    print("========= file written to ",args.output_path),
    print("Done!")